package leetcode_1501_1600;

import javax.xml.stream.events.EndElement;
import java.util.*;

public class LeeCode_1514 {
    public static void main(String[] args) {
        System.out.println(maxProbability(3, new int[][]{{0, 1}, {1, 2}, {0, 2}}, new double[]{0.5, 0.5, 0.2}, 0, 2));
    }
    static class Node{
        int i;
        double p;

        public Node(int i, double p) {
            this.i = i;
            this.p = p;
        }
    }
    private static double maxProbability(int n, int[][] edges, double[] succProb, int start, int end) {
        Map<Integer, List<int[]>> map = new HashMap<>();
        for (int i = 0; i < edges.length; i++) {
            int[] edge = edges[i];
            if (map.containsKey(edge[0])){
                map.get(edge[0]).add(new int[]{edge[1], i});
            }else {
                List<int[]> list = new ArrayList<>();
                list.add(new int[]{edge[1], i});
                map.put(edge[0], list);
            }
            if (map.containsKey(edge[1])){
                map.get(edge[1]).add(new int[]{edge[0], i});
            }else {
                List<int[]> list = new ArrayList<>();
                list.add(new int[]{edge[0], i});
                map.put(edge[1], list);
            }
        }
        Map<Integer, Double> shortest = new HashMap<>();
        shortest.put(start, 1.0);
        PriorityQueue<Node> queue = new PriorityQueue<>((o1, o2) -> {
            if (o1.p == o2.p)
                return 0;
            return o2.p - o1.p > 0 ? 1 : -1;
        });
        if (map.containsKey(start)){
            for (int[] ints : map.get(start)) {
                queue.add(new Node(ints[0], succProb[ints[1]]));
            }
        }
        while (!queue.isEmpty()){
            Node poll = queue.poll();
            int cur = poll.i;
            if (shortest.containsKey(cur))
                continue;
            double p = poll.p;
            shortest.put(cur, p);
            for (int[] next : map.get(cur)) {
                if (shortest.containsKey(next[0])){
                    continue;
                }
                queue.add(new Node(next[0], succProb[next[1]] * p));
            }
        }
        return shortest.getOrDefault(end, 0.0);
    }
}
